# -*- coding: utf-8 -*-

from mitmproxy.options import Options
from mitmproxy.http import HTTPFlow
from mitmproxy.connection import Client
from mitmproxy.addons.proxyserver import Proxyserver
from mitmproxy.tools.dump import DumpMaster
from mitmproxy.http import cookies as http_cookies

import json
import base64

from AddonBase import AddonBase
import asyncio
from ZiMoooModule.TcpPullSeverClient import Client
import time


def GetCookieExpires(timeStamp=None, MaxAge=None):
    x = {1: 'Mon', 2: 'Tue', 3: 'Wed', 4: 'Thu', 5: 'Fri', 6: 'Sat', 7: 'Sun'}
    m = {1: "Jan", 2: "Feb", 3: "Mar", 4: "Apr", 5: "May", 6: "Jun", 7: "Jul", 8: "Aug", 9: "Sept", 10: "Oct",
         11: "Nov", 12: "Dec"}
    if not timeStamp:
        timeStamp = time.time()
    if not MaxAge:
        MaxAge = 365 * 24 * 60 * 60
    timeStamp = timeStamp + MaxAge
    timeArray = time.localtime(timeStamp)
    s = time.strftime(", %d-YUE-%Y %H:%M:%S GMT", timeArray)
    s = x[timeArray.tm_wday + 1] + s
    s = s.replace("YUE", m[timeArray.tm_mon])
    return s


class MyAddon(AddonBase):
    def __init__(self, SocketClient=None, socket_ip=None, socket_port=None, CONFIGPATH=None):
        super().__init__()
        try:
            with open(CONFIGPATH, "r", encoding="utf-8") as fp:
                self.CONFIG = json.load(fp)
        except Exception as E:
            print(E)
        self.client = None
        self.socket_ip = socket_ip
        self.socket_port = socket_port
        if self.socket_ip:
            self.client = SocketClient(self.socket_ip, self.socket_port)
            self.client.Start()
        self.REQUEST_LIST = []
        pass

    def FlowDataToJson(self, flow: HTTPFlow, rtype="request"):
        if rtype == "request":
            re = flow.request
        else:
            re = flow.response
        resdata = {}
        resdata['urlencoded_form'] = []
        resdata['query'] = []
        if rtype == "request":
            resdata['Url'] = re.url
            resdata['Method'] = re.method
            print("urlencoded_form", re.urlencoded_form)
            print("query", re.query)
            if re.urlencoded_form:
                for key in re.urlencoded_form:
                    value = re.urlencoded_form.get(key)
                    print("urlencoded_form-value-type", type(value))
                    resdata['urlencoded_form'].append({"name": key, "value": value})
                # try:
                #     post_p = eval(reduce(lambda x, y: x + y, flow.request.urlencoded_form.keys()))
                #     post_param = {k: v for k, v in map(lambda x: x.split(":"), post_p[1:-1].split(","))}
                # except Exception:
                #     pass
            if re.query:
                for key in re.query:
                    value = re.query.get(key)
                    print("query-value-type", type(value))
                    resdata['query'].append({"name": key, "value": value})
        cookies = []
        for key in re.cookies:
            # print("cookies-value", re.cookies.get(key))
            # a=('1', CookieAttrs[('expires', 'Fri, 31-May-2024 07:02:16 GMT'), ('Max-Age', '31536000'), ('path', '/'), (
            # 'domain', '.125.la'), ('secure', '')])
            if rtype == "response":
                value = re.cookies.get(key)
                cookies.append({"name": key, "value": value[0], "domain": value[1].get("domain")})
            else:
                cookies.append({key: re.cookies.get(key)})

        resdata['Cookies'] = cookies
        headers = []
        for key in re.headers:
            if key == "Cookie":
                pass
            elif key == "CurSocketConnID":
                pass
            else:
                headers.append({"name": key, "value": re.headers.get(key)})
        resdata['Headers'] = headers
        if rtype == "response":
            resdata['Code'] = re.status_code
            resdata['Data'] = base64.b64encode(re.content).decode()
        elif resdata['Method'] == "POST":
            resdata['Data'] = base64.b64encode(re.content).decode()
            print("POSTDATA", re.content.decode())
        return resdata

    def request(self, flow: HTTPFlow):
        try:
            BrowserID = int(flow.request.headers.pop("____BrowserID____"))
            capture_type = int(flow.request.headers.pop("____capture_type____"))
            QueryID = flow.request.headers.pop("____QueryID____")
            proxy_type = int(flow.request.headers.pop("____proxy_type____"))
            proxy_user = flow.request.headers.pop("____proxy_user____")
            proxy_server = flow.request.headers.pop("____proxy_server____")
            ProcessID = flow.request.headers.pop("____ProcessID____")
            self.REQUEST_LIST.append({"FlowID": flow.id,
                                      'BrowserID': BrowserID,
                                      "capture_type": capture_type,
                                      "QueryID": QueryID,
                                      "ProcessID": ProcessID
                                      })
        except Exception as E:
            print("非浏览器代理包", E)
            return
        # requstData = self.FlowDataToJson(flow=flow, rtype="request")

        if self.CONFIG["CAPTURE_TYPE"]["REQUEST_MODIFY"] == self.CONFIG["CAPTURE_TYPE"][
            "REQUEST_MODIFY"] & capture_type:
            requstData = self.FlowDataToJson(flow=flow, rtype="request")
            requstData["QueryID"] = QueryID
            requstData["BrowserID"] = BrowserID
            requstData["ProcessID"] = ProcessID
            jsondata = json.dumps(requstData)

            try:
                ServerResponse = self.client.Sync_SendData(apiCode=self.CONFIG["EVENT"]["ON_REQUEST_MODIFY"],
                                                           data=jsondata.encode(),
                                                           outTime=15)
                print("修改请求", ServerResponse)
                ServerResponse = json.loads(ServerResponse)
                NewUrl = ServerResponse["data"]["NewUrl"]
                NewData = ServerResponse["data"]["NewData"]
                NewHeaders = ServerResponse["data"]["NewHeaders"]
                NewCookies = ServerResponse["data"]["NewCookies"]
                NewUrlencodedForms=ServerResponse["data"]["NewUrlencodedForms"]
                NewQuerys = ServerResponse["data"]["NewQuerys"]

                if NewUrl:
                    flow.request.url = NewUrl
                elif NewQuerys:
                    flow.request.query.clear()
                    for item in NewQuerys:
                        flow.request.query.add(key=item["name"],value=item["value"])

                if NewData:
                    flow.request.content = base64.b64decode(NewData)
                elif NewUrlencodedForms:
                    flow.request.urlencoded_form.clear()
                    for item in NewUrlencodedForms:
                        flow.request.urlencoded_form.add(key=item["name"],value=item["value"])

                if NewHeaders:
                    for item in NewHeaders:
                        flow.request.headers[item["name"]] = item["value"]

                if NewCookies:
                    for item in NewCookies:
                        try:
                            flow.request.cookies.add(item["name"], (item["value"]))
                        except Exception as E:
                            print("Cookies设置异常", E)
                    pass



            except Exception as E:
                print("修改请求：返回数据异常", E)
                # flow.request.query.set_all("type", ["2"])
                # flow.request.urlencoded_form = [
                #     ("keyWords", "米"), ("channelId", "28"), ("cityId", "306")
                # ]
            # print("修改请求：用时", time.time() - t)

        # 以下为代理设置
        # print("设置代理","proxy_type",proxy_type,"proxy_server", proxy_server, "username", proxy_user)
        if proxy_type == self.CONFIG["PROXY_TYPE"]["NONE"]:
            self.CanelProxy(flow)
            return

        if proxy_type == self.CONFIG["PROXY_TYPE"]["HTTP"]:
            proxy_server = "http://" + proxy_server
        elif proxy_type == self.CONFIG["PROXY_TYPE"]["HTTPS"]:
            proxy_server = "https://" + proxy_server
        elif proxy_type == self.CONFIG["PROXY_TYPE"]["SOCKET"]:
            proxy_server = "socket://" + proxy_server
        elif proxy_type == self.CONFIG["PROXY_TYPE"]["SOCKET4"]:
            proxy_server = "socket4://" + proxy_server
        elif proxy_type == self.CONFIG["PROXY_TYPE"]["SOCKET5"]:
            proxy_server = "socket5://" + proxy_server
        else:
            proxy_server = "http://" + proxy_server
        # print("设置代理：proxy_server",proxy_server,"username",proxy_user)
        self.SetProxy(flow, proxy_server, username=proxy_user)

    def response(self, flow: HTTPFlow):
        BrowserID = None
        QueryID = None
        ProcessID = None
        capture_type = self.CONFIG["CAPTURE_TYPE"]["NONE"]
        self.FlowDataToJson(flow=flow, rtype="response")
        for i in range(len(self.REQUEST_LIST)):
            if self.REQUEST_LIST[i]["FlowID"] == flow.id:
                item = self.REQUEST_LIST.pop(i)
                BrowserID = item["BrowserID"]
                QueryID = item["QueryID"]
                capture_type = item["capture_type"]
                ProcessID = item["ProcessID"]
                break
        if not (BrowserID and QueryID):
            return

        if self.CONFIG["CAPTURE_TYPE"]["RESPONSE_MODIFY"] == self.CONFIG["CAPTURE_TYPE"][
            "RESPONSE_MODIFY"] & capture_type:
            responseData = self.FlowDataToJson(flow=flow, rtype="response")
            responseData["QueryID"] = QueryID
            responseData["BrowserID"] = BrowserID
            responseData["ProcessID"] = ProcessID
            jsondata = json.dumps(responseData)
            try:
                ServerResponse = self.client.Sync_SendData(apiCode=self.CONFIG["EVENT"]["ON_RESPONSE_MODIFY"],
                                                           data=jsondata.encode(),
                                                           outTime=15)
                # print("ServerResponse",ServerResponse)
                ServerResponse = json.loads(ServerResponse)
                # print("ServerResponse",ServerResponse)
                if ServerResponse["data"]:
                    NewData = ServerResponse["data"]["NewData"]
                    NewHeaders = ServerResponse["data"]["NewHeaders"]
                    NewCookies = ServerResponse["data"]["NewCookies"]
                    if NewData:
                        flow.response.content = base64.b64decode(NewData)
                    if NewHeaders:
                        for item in NewHeaders:
                            flow.response.headers[item["name"]] = item["value"]


                    if NewCookies:
                        for item in NewCookies:
                            CookieAttrs = http_cookies.CookieAttrs()
                            # CT.add("aa", "bb")
                            # cookies = []
                            # print("cookies-value", re.cookies.get(key))
                            # 'Thu Jun  1 15:30:07 2023 GMT'
                            # a=('1', CookieAttrs[('expires', 'Fri, 31-May-2024 07:02:16 GMT'), ('Max-Age', '31536000'), ('path', '/'), (
                            # 'domain', '.125.la'), ('secure', '')])

                            CookieAttrs.add("domain", item["domain"])
                            MaxAge = 365 * 24 * 60 * 60
                            CookieAttrs.add("expires", GetCookieExpires(timeStamp=time.time(), MaxAge=MaxAge))
                            CookieAttrs.add("Max-Age", str(MaxAge))
                            CookieAttrs.add("path", "/")
                            CookieAttrs.add("secure", "")
                            try:
                                flow.response.cookies.add(item["name"], (item["value"], CookieAttrs))
                            except Exception as E:
                                print("Cookies设置异常", E)
                        pass

            except Exception as E:
                print("修改响应：返回数据异常", E)

        # if self.CONFIG["CAPTURE_TYPE"]["RESPONSE_READ"] == self.CONFIG["CAPTURE_TYPE"]["RESPONSE_READ"] & capture_type:
        #     responseData = self.FlowDataToJson(flow=flow, rtype="response")
        #     responseData["QueryID"] = QueryID
        #     responseData["BrowserID"] = BrowserID
        #     responseData["ProcessID"] = ProcessID
        #     jsondata = json.dumps(responseData)
        #
        #     self.client.SendData(apiCode=self.CONFIG["EVENT"]["ON_RESPONSE"], data=jsondata.encode(),
        #                          outTime=30)
        # # print("修改响应协议",ServerResponse)


def MitmProxyStart(socket_ip=None, socket_port=None, CONFIGPATH=None):
    async def _main():
        # 初始设置代理模式 默认：regular
        # mode: Sequence[str] = ["upstream:http://localhost:7889"]
        # options = Options(listen_host='localhost', listen_port=8889, mode=mode)
        options = Options(listen_host='localhost', listen_port=8889)
        options.add_option("connection_strategy", str, "lazy", "")
        options.add_option("upstream_cert", bool, False, "")
        options.add_option("ssl_insecure", bool, True, "")
        # EventLog verbosity. Default: info  Choices: error, warn, info, alert, debug
        # 事件日志 默认为 info
        options.add_option("console_eventlog_verbosity", str, "error", "")
        # Log verbosity.. Default: info  Choices: error, warn, info, alert, debug
        options.add_option("termlog_verbosity", str, "error", "")
        # 替换代理模式
        # mode: Sequence[str] = ["upstream:http://localhost:7889"]
        # mode_option = {'mode': mode}
        # options.update(**mode_option)
        m = DumpMaster(options=options)
        m.server = Proxyserver()
        addons = MyAddon(SocketClient=Client, socket_ip=socket_ip, socket_port=socket_port, CONFIGPATH=CONFIGPATH)
        m.addons.add(addons)
        await m.run()

    try:
        asyncio.run(_main())
    except KeyboardInterrupt:
        pass


if __name__ == '__main__':
    import sys

    print('参数个数为:', len(sys.argv), '个参数。')
    print('参数列表:', str(sys.argv))
    print('脚本名为：', sys.argv[0])
    CONFIGPATH = ""
    for item in sys.argv:
        print(item)
        print("item[0:11]", item[0:11])
        if len(item) > 12 and item[0:11] == "CONFIGPATH=":
            CONFIGPATH = item[11:]

    print("CONFIGPASTH", CONFIGPATH)

    MitmProxyStart(socket_ip="127.0.0.1", socket_port=7889, CONFIGPATH=CONFIGPATH)
    # 注释掉 C:\soft\python39\Lib\site-packages\mitmproxy\addons\dumper.py
    # 244行 和 253行 输出的日志
    '''
    pyinstaller pymitmdump\mitmdumpStartApi.py --distpath pymitmdumpstartdist --add-data "pymitmdump\mitmdump_executable;mitmdump_executable" --add-data "pymitmdump\mitmdumpUrlSaver.py;."
    pyinstaller pymitmdump\mitmdumpOtherApi.py --distpath pymitmdumpotherdist
    '''
